import pymysql
import socket
import datetime
import sys
import ipaddress
import threading
import os
from utils import *
from scapy.sendrecv import sr
from scapy.layers.inet import IP, TCP


BLUE, RED, WHITE, YELLOW, MAGENTA, GREEN, END = '', '', '', '', '', '', ''

class ThreadManager(object):
    i = 0

    def __init__(self, ipList):
        self.allIps = ipList
        self.size = len(ipList)

    def getNextIp(self):
        if not (self.i >= self.size - 1):
            ip = self.allIps[self.i]
            self.i += 1
            return ip
        return 0

    def getID(self):
        return self.i + 1

def checkSQL(host, port):
    passwds = ['', 'admin', '12345678', '11111111']
    TruePassWord = None
    loginFail = False
    try:
        con = pymysql.connect(host=host, connect_timeout=sqlTimeout)
    except:
        loginFail = True

    if loginFail:
        for pw in passwds:
            try:
                con = pymysql.connect(host=host, user='root', passwd=pw, connect_timeout=sqlTimeout)
                TruePassWord = pw
                break
            except:
                pass
        else:
            return ["login-error"]

    try:
        cursor = con.cursor()
        cursor.execute("SHOW DATABASES")
        dbs = cursor.fetchall()
    except:
        dbs = ["-"]

    try:
        cursor = con.cursor()
        cursor.execute("SELECT VERSION()")
        serverVersion = str(cursor.fetchone())
    except:
        serverVersion = "-"

    return ["success", serverVersion, dbs, TruePassWord]


def coreOptions():
    options = [["network", "IP range to scan", ""],
               ["port", "Port to scan.", "3306"],
               ["port-timeout", "Timeout (in sec) for port 80.", "0.3"],
               ["sql-timeout", "Timeout (in sec) for the database connection.", "3"],
               ["threads", "Number of threads to run.", "50"],
               ["checkauth", "Connect to the server and perform tests.", "true"],
               ["verbose", "Show verbose output.", "true"]]
    return options


def createIPList(network):
    net4 = ipaddress.ip_network(network)
    ipList = []
    for x in net4.hosts():
        ipList.append(x)
    return ipList


def print1(data):
    global ui
    if verbose:
        print(data)
        ui.logtext.append(data)

def print2(data):
    global ui
    if verbose:
        print(data)
        ui.logtext.append(data)
        # ui.result.append(data)

# def checkServer(address, port):
#     #检查的逻辑就是去尝试着连接一下，看结果如何
#     #顺便考虑一下是使用SYN连接还是直接connet
#     global method
#     if method == "connet":
#         s = socket.socket()
#         s.settimeout(float(portTimeout))
#         try:
#             s.connect((address, port))
#             s.close()
#             return True
#         except socket.error:
#             s.close()
#             return False
#         except:
#             s.close()
#             return "FAIL"
#     elif method == "syn":
#         try:
#             temp = sr(IP(dst = address)/
#                       TCP(dport=(int(port)), flags='S'),
#                       timeout=2, verbose=False)
#             if temp[0].res:
#                 result = temp[0].res  # temp分回复和无回显
#                 if (result[0][1].payload.flags) == 'SA':
#                     return True
#                 else:
#                     return False
#             else:
#                 return False
#         except:
#             return 'FAIL'


def writeToFile(line):
    file = open(fileName, "a")
    file.write(line)
    ui.result.append(line)
    file.close()


def restart_line():
    sys.stdout.write('\r')
    sys.stdout.flush()


def statusWidget():
    sys.stdout.write(GREEN + "[" + status + "] " + YELLOW + str(threadManager.getID()) + GREEN + " / " + YELLOW + str(
        allIPs) + GREEN + " hosts done." + END)
    restart_line()
    sys.stdout.flush()


def scan(i):
    global status
    global openPorts
    global done
    while True:
        if stop:
            sys.exit()
        ip = threadManager.getNextIp()
        if ip == 0:
            break
        status = (threadManager.getID() / allIPs) * 100
        status = format(round(status, 2))
        status = str(status) + "%"
        stringIP = str(ip)
        isUp = checkServer(stringIP, port, portTimeout, method)
        if isUp != "FAIL":
            if isUp:
                openPorts = openPorts + 1
                print2("[+] Port " + str(port) + " is open on '" + stringIP + "'" + END)

                if checkauth:
                    mysqlStatusReason = "UNKNOWN ERROR"
                    mysql = checkSQL(stringIP, port)
                    if mysql[0] == "login-error":
                        mysqlStatus = False
                        mysqlStatusReason = "CONNECTION ERROR"
                        print1(
                            "[!] Failed connecting to the database on '" + stringIP + "'. ERROR: " + mysqlStatusReason + END)
                    elif mysql[0] == "permission-error":
                        mysqlStatus = False
                        mysqlStatusReason = "PERMISSION ERROR"
                        print1(
                            "[!] Failed connecting to the database on '" + stringIP + "'. ERROR: " + mysqlStatusReason + END)
                    elif mysql[0] == "success":
                        mysqlStatus = True
                        version = mysql[1]
                        dbs = mysql[2]
                        passWord = mysql[3]
                        dbsList = ""
                        for db in dbs:
                            dbsList = dbsList + str(db) + ", "
                        if dbsList != "":
                            dbsList = dbsList[:-2]
                        else:
                            dbsList = "-"

                        print2("[+] Open database found:\n\tIP: " + stringIP + "\n\t" + "MySQL version: " + str(
                            version) + "\n\tDB's: " + dbsList + "\nPassword of root is: " + passWord + "\n")

                    else:
                        print1(
                            "[!] Failed connecting to the database on '" + stringIP + "'. ERROR: " + mysqlStatusReason)
                        mysqlStatus = False

                    if mysqlStatus:
                        logLine = stringIP + " - " + str(
                            port) + " OPEN" + " - " + "OPEN DATABASE - Version: " + version + " - " + " DB's: " + dbsList + "\n"

                    else:
                        logLine = stringIP + " - " + str(
                            port) + " OPEN" + " - DB SCAN ERROR: " + mysqlStatusReason + "\n"
                    print2(logLine)
                    logLines.append(logLine)
                else:
                    logLine = stringIP + " - " + str(port) + " OPEN\n"
                    logLines.append(logLine)
            elif not isUp:
                print1("[-] Port " + str(port) + " is closed on '" + stringIP + "'")
        else:
            print1("[!] Failed connecting to '" + stringIP + "'")
    done = done + 1


def core(moduleOptions, ui_ori):
    print(
        "\n" + GREEN + "MYSQL module by @xdavidhu. Scanning subnet '" + YELLOW + moduleOptions[0][2] + GREEN + "'...\n")

    global status
    global fileName
    global allIPs
    global portTimeout
    global titleTimeout
    global ips
    global threadCount
    global done
    global verbose
    global stop
    global port
    global openPorts
    global logLines
    global checkauth
    global sqlTimeout
    global threadManager
    global ui
    global method
    ui = ui_ori
    logLines = []
    stop = False
    done = 0
    try:
        port = int(moduleOptions[1][2])
    except:
        print(RED + "[!] Invalid port. Exiting...\n")
        return
    portTimeout = moduleOptions[2][2]
    network = moduleOptions[0][2]
    sqlTimeout = moduleOptions[3][2]
    threadCount = int(moduleOptions[4][2])
    checkauth = moduleOptions[5][2]
    verbose = moduleOptions[6][2]

    try:
        sqlTimeout = int(sqlTimeout)
    except:
        print(RED + "[!] Invalid sql-timeout. Exiting...\n")
        return

    if verbose == "true":
        verbose = True
    else:
        verbose = False

    if checkauth == "true":
        checkauth = True
    else:
        checkauth = False

    try:
        ipList = createIPList(network)
        allIPs = len(ipList)
        if allIPs == 0:
            raise Exception
    except:
        print(RED + "[!] Invalid subnet. Exiting...\n")
        return

    method = 'connect' if ui.connect.isChecked() else 'syn'

    threadManager = ThreadManager(ipList)

    i = datetime.datetime.now()
    i = str(i).replace(" ", "_")
    i = str(i).replace(":", "-")
    script_path = os.path.dirname(os.path.realpath(__file__))
    script_path = script_path.replace("modules", "")
    if not os.path.exists(script_path + "logs"):
        os.makedirs(script_path + "logs")
    fileName = script_path + "logs/log-mysql-portSpider-" + i + ".log"

    file = open(fileName, 'w')
    file.write("subnet: " + network + "\n")
    file.close()

    openPorts = 0
    threads = []
    for i in range(threadCount):
        t = threading.Thread(target=scan, args=(i,))
        t.daemon = True
        threads.append(t)
        t.start()

    try:
        while True:
            if done == threadCount and threadManager.getID() == allIPs:
                break
            statusWidget()
    except KeyboardInterrupt:
        stop = True
        verbose = False
        print("\n" + RED + "[I] Stopping..." + END)
    stop = True
    verbose = False

    for logLine in logLines:
        try:
            writeToFile(logLine)
        except:
            writeToFile("WRITING-ERROR")
    print("\n\n" + "[I] MYSQL module done. Results saved to '" + fileName + "'.\n")
